# Pardiso.jl Schur Complement Matrix Processing Example
# Author: MiniMax Agent
# This example demonstrates how to use Pardiso.jl for efficient Schur complement computation

using Pardiso
using SparseArrays
using LinearAlgebra
using Random

# Set random seed for reproducibility
Random.seed!(123)

"""
    compute_schur_complement_pardiso(A, B, C, D)

Compute the Schur complement S = D - C * A^(-1) * B using Pardiso solver.

For a block matrix:
M = [A  B]
    [C  D]

The Schur complement of A in M is: S = D - C * A^(-1) * B

Parameters:
- A: Upper-left block (must be invertible)
- B: Upper-right block  
- C: Lower-left block
- D: Lower-right block

Returns:
- S: Schur complement matrix
"""
function compute_schur_complement_pardiso(A, B, C, D)
    println("Computing Schur complement using Pardiso.jl...")
    
    # Initialize Pardiso solver
    ps = PardisoSolver()
    
    # Set matrix type (11 = real unsymmetric matrix)
    set_matrixtype!(ps, Pardiso.REAL_UNSYM)
    
    # Factor matrix A
    pardisoinit(ps)
    set_phase!(ps, Pardiso.ANALYSIS_NUM_FACT)
    
    # Convert A to the required format if needed
    A_sparse = sparse(A)
    
    println("Factorizing matrix A...")
    pardiso(ps, A_sparse, B)
    
    # Solve A * X = B to get X = A^(-1) * B
    set_phase!(ps, Pardiso.SOLVE_ITERATIVE_REFINE)
    X = similar(B)
    pardiso(ps, X, A_sparse, B)
    
    # Compute Schur complement: S = D - C * X
    S = D - C * X
    
    println("Schur complement computed successfully!")
    
    # Clean up Pardiso memory
    set_phase!(ps, Pardiso.RELEASE_ALL)
    pardiso(ps)
    
    return S
end

"""
    generate_test_matrices(n1, n2)

Generate test matrices for Schur complement computation.

Parameters:
- n1: Size of block A (n1 × n1)
- n2: Size of block D (n2 × n2)

Returns:
- A, B, C, D: Block matrices
"""
function generate_test_matrices(n1, n2)
    println("Generating test matrices...")
    
    # Create a well-conditioned matrix A
    A = randn(n1, n1)
    A = A' * A + I  # Make it positive definite
    
    # Create other blocks
    B = randn(n1, n2)
    C = randn(n2, n1)
    D = randn(n2, n2)
    D = D' * D + I  # Make it positive definite
    
    return A, B, C, D
end

"""
    verify_schur_complement(A, B, C, D, S)

Verify the Schur complement computation by checking the block matrix properties.
"""
function verify_schur_complement(A, B, C, D, S)
    println("Verifying Schur complement computation...")
    
    # Direct computation for verification
    A_inv = inv(A)
    S_direct = D - C * A_inv * B
    
    # Check if results match
    error = norm(S - S_direct)
    println("Verification error: $(error)")
    
    if error < 1e-10
        println("✓ Schur complement computation verified!")
    else
        println("✗ Verification failed!")
    end
    
    return error < 1e-10
end

"""
    demonstrate_block_matrix_inversion(A, B, C, D, S)

Demonstrate how Schur complement can be used for block matrix inversion.
"""
function demonstrate_block_matrix_inversion(A, B, C, D, S)
    println("\nDemonstrating block matrix inversion using Schur complement...")
    
    # For block matrix M = [A B; C D], the inverse is:
    # M^(-1) = [A^(-1) + A^(-1)*B*S^(-1)*C*A^(-1)  -A^(-1)*B*S^(-1)]
    #          [-S^(-1)*C*A^(-1)                      S^(-1)        ]
    
    A_inv = inv(A)
    S_inv = inv(S)
    
    # Compute blocks of the inverse
    M11 = A_inv + A_inv * B * S_inv * C * A_inv
    M12 = -A_inv * B * S_inv
    M21 = -S_inv * C * A_inv
    M22 = S_inv
    
    # Construct the inverse matrix
    M_inv = [M11 M12; M21 M22]
    
    # Original block matrix
    M = [A B; C D]
    
    # Verify inversion
    I_check = M * M_inv
    error = norm(I_check - I)
    println("Block matrix inversion error: $(error)")
    
    return M_inv
end

"""
    main_example()

Main function demonstrating various aspects of Schur complement computation with Pardiso.jl
"""
function main_example()
    println("=== Pardiso.jl Schur Complement Matrix Processing Example ===\n")
    
    # Matrix dimensions
    n1, n2 = 50, 30
    
    # Generate test matrices
    A, B, C, D = generate_test_matrices(n1, n2)
    
    println("Matrix dimensions:")
    println("A: $(size(A)), B: $(size(B))")
    println("C: $(size(C)), D: $(size(D))")
    println()
    
    # Compute Schur complement using Pardiso
    S = compute_schur_complement_pardiso(A, B, C, D)
    
    println("Schur complement S dimensions: $(size(S))")
    println()
    
    # Verify the computation
    verified = verify_schur_complement(A, B, C, D, S)
    
    # Demonstrate block matrix inversion
    if verified
        M_inv = demonstrate_block_matrix_inversion(A, B, C, D, S)
        println("✓ Block matrix inversion completed successfully!")
    end
    
    println("\n=== Example completed ===")
end

# Additional utility functions for advanced use cases

"""
    solve_schur_system(A, B, C, D, f, g)

Solve a block linear system using Schur complement method.

System: [A B] [x] = [f]
        [C D] [y]   [g]

Solution approach:
1. Compute Schur complement S = D - C * A^(-1) * B
2. Solve S * y = g - C * A^(-1) * f
3. Solve A * x = f - B * y
"""
function solve_schur_system(A, B, C, D, f, g)
    println("Solving block system using Schur complement method...")
    
    # Step 1: Compute Schur complement
    S = compute_schur_complement_pardiso(A, B, C, D)
    
    # Step 2: Solve for y
    A_inv_f = A \ f
    rhs_y = g - C * A_inv_f
    y = S \ rhs_y
    
    # Step 3: Solve for x
    x = A \ (f - B * y)
    
    return x, y
end

"""
    performance_comparison(sizes)

Compare performance of different Schur complement computation methods.
"""
function performance_comparison(sizes)
    println("Performance comparison for different matrix sizes:")
    println("Size\tPardiso Time\tDirect Time\tSpeedup")
    println("-" ^ 50)
    
    for n in sizes
        A, B, C, D = generate_test_matrices(n, n)
        
        # Time Pardiso method
        pardiso_time = @elapsed begin
            S_pardiso = compute_schur_complement_pardiso(A, B, C, D)
        end
        
        # Time direct method
        direct_time = @elapsed begin
            S_direct = D - C * (A \ B)
        end
        
        speedup = direct_time / pardiso_time
        println("$(n)\t$(round(pardiso_time, digits=4))\t\t$(round(direct_time, digits=4))\t\t$(round(speedup, digits=2))x")
    end
end

# Run the main example
if abspath(PROGRAM_FILE) == @__FILE__
    main_example()
    
    # Uncomment to run performance comparison
    # performance_comparison([10, 20, 50, 100])
end